#############################################################################

#     Weighted imputations method for mediation analysis               
#     Applicable situation:
#         1.The outcome is dichotomous data;
#         2.An exposure variable and two covariates: dichotomous data;
#         2.Two parallel mediators.

#############################################################################
WeightedImputation<-function(Data,n){
  Data$id<-c(1:n)
    
  #Fit a suitable model for the probability (density) of the mediators
  if(all(unique(Data$M1)%in%c(0,1))){
    fitM1 <- glm(M1 ~ A+C1+C2,family = binomial("logit"), data = Data)
  }else{
    fitM1 <- glm(M1 ~ A+C1+C2,family = gaussian("identity"), data = Data)
  }
  
  if(all(unique(Data$M2)%in%c(0,1))){
    fitM2 <- glm(M2 ~ A+C1+C2,family = binomial("logit"), data = Data)
  }else{
    fitM2 <- glm(M2 ~ A+C1+C2,family = gaussian("identity"), data = Data)
  }
  
  #Fit a suitable model for the outcome mean
  fitY <- glm(Y ~ A*M1*M2+C1+C2,family = binomial("logit"), data = Data)
  
  #Construct the extended data set
  extdat <- data.frame(replicate = rep(1:4, times = nrow(Data)),Data[rep(Data$id, each = 4), ],a0 = NA, a1 = NA, a2 = NA)
  
  extdat1<-within(extdat,{
    a0 <- ifelse(replicate %in% c(2,4), 1-A, A)
    a1 <- ifelse(replicate %in% c(3,4), 1-A, A)
    a2<-A
  })
  
  extdat2<-within(extdat,{
    a0 <- ifelse(replicate %in% c(2,4), 1-A, A)
    a1<-A
    a2 <- ifelse(replicate %in% c(3,4), 1-A, A)
  })
  
  #Calculate regression weights
  #calculate W1 for mediator M1
  if(all(unique(Data$M1)%in%c(0,1))){
    num1 <- with(extdat1,
                 dbinom(M1, size = 1, prob = predict(fitM1, newdata = within(extdat1, A <- a1),type = "response")))
    
    denom1 <- with(extdat1,
                   dbinom(M1, size = 1, prob = predict(fitM1, newdata = within(extdat1, A <- a2),type = "response")))
    
    extdat1$W1 <- num1/denom1
  }else{
    num1 <- with(extdat1,
                 dnorm(M1, mean = predict(fitM1, newdata = within(extdat1, A <- a1),type = "response"),
                       sd = sqrt(summary(fitM1)$dispersion)))
    
    denom1 <- with(extdat1,
                   dnorm(M1, mean = predict(fitM1, newdata = within(extdat1, A <- a2),type = "response"),
                         sd = sqrt(summary(fitM1)$dispersion)))
    
    extdat1$W1 <- num1/denom1
  }
  
  #calculate W2 for mediator M2
  if(all(unique(Data$M2)%in%c(0,1))){
    num2 <- with(extdat2,
                 dbinom(M2, size = 1, prob = predict(fitM2, newdata = within(extdat2, A <- a2),type = "response")))
    
    denom2 <- with(extdat2,
                   dbinom(M2, size = 1, prob = predict(fitM2, newdata = within(extdat2, A <- a1),type = "response")))
    
    extdat2$W2 <- num2/denom2
  }else{
    num2 <- with(extdat2,
                 dnorm(M2, mean = predict(fitM2, newdata = within(extdat2, A <- a2),type = "response"),
                       sd = sqrt(summary(fitM2)$dispersion)))
    
    denom2 <- with(extdat2,
                   dnorm(M2, mean = predict(fitM2, newdata = within(extdat2, A <- a1),type = "response"),
                         sd = sqrt(summary(fitM2)$dispersion)))
    
    extdat2$W2 <- num2/denom2
  }
  
  #Compare regression weights and select the working model(M1 or M2)
  library(dplyr)
  diff<-extdat1$W1-extdat2$W2
  flag<-case_when(diff>0 ~ 1,
                  diff==0 ~ 0,
                  diff<0 ~(-1))
  if(sum(flag)>=0){
    #Impute nested counterfactuals 
    extdat1$Y <- predict(fitY, newdata = within(extdat1, A <- a0),type = "response")
    
    #Fit a natural effect model
    fitNEM <- glm(Y ~ a0 * a1 * a2 + C1 + C2,
                   family = binomial("logit"), data = extdat1, weights = W1)
  }else{
    #Impute nested counterfactuals
    extdat2$Y <- predict(fitY, newdata = within(extdat2, A <- a0),type = "response")
    
    #Fit a natural effect model
    fitNEM <- glm(Y ~ a0 * a1 * a2 + C1 + C2,
                   family = binomial("logit"), data = extdat2, weights = W2)
  }
  
  # ###Obtain population-average component effects(rather than effects conditional on C)
  # fitA <- glm(A ~ C1+C2,family = binomial("logit"), data = Data)
  # #updated weights
  # extdat1 <- within(extdat1, {
  #   W1 <- W1 / dbinom(A, size = 1, prob = predict(fitA, newdata = extdat1,type = "response"))})
  # extdat2 <- within(extdat2, {
  #   W2 <- W2 / dbinom(A, size = 1, prob = predict(fitA, newdata = extdat2,type = "response"))})
  # #Fit a population-average natural effect model
  # fitNEM1 <- glm(Y ~ a0 * a1 * a2,
  #                family = binomial("logit"), data = extdat1, weights = W1)
  # fitNEM2 <- glm(Y ~ a0 * a1 * a2,
  #                family = binomial("logit"), data = extdat2, weights = W2)
  
  return(coef(fitNEM))
}